import pandas as pd
import torch
import torch.nn as nn
from sklearn import model_selection

df = pd.read_excel('B22.xlsx')

class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(14, 100),
            nn.ReLU(),
            nn.Linear(100, 30),
            nn.ReLU(),
            nn.Linear(30, 2)
        )

    def forward(self, x):
        logits = self.linear_relu_stack(x)
        return logits

device = 'cuda'

data = df.iloc[:, 6:20].to_numpy()
labels = df.iloc[:, 4].to_numpy()

kf = model_selection.KFold(n_splits=6, shuffle=True)
results = []
for X_train_i,X_test_i in kf.split(data):
  print(f'Train size: {len(X_train_i)}')
  print(f'Test size: {len(X_test_i)}')
  model = NeuralNetwork().to(device)
  loss_fn = nn.CrossEntropyLoss()
  optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

  model.train()
  for epoch in range(20):
      X = torch.from_numpy(data[X_train_i]).float().to(device)
      y = torch.from_numpy(labels[X_train_i]).long().to(device)

      # Compute prediction error
      pred = model(X)
      loss = loss_fn(pred, y)

      # Backpropagation
      optimizer.zero_grad()
      loss.backward()
      optimizer.step()

      accuracy = torch.count_nonzero(pred.argmax(axis=-1) == y) / len(X_train_i)

      print(f'Epoch: {epoch + 1}\tLoss:{loss.item():.4f}\tAccuracy:{accuracy:.3f}')
    
  model.eval()

  X = torch.from_numpy(data[X_test_i]).float().to(device)
  y = torch.from_numpy(labels[X_test_i]).long().to(device)

  pred = model(X)
  accuracy = torch.count_nonzero(pred.argmax(axis=-1) == y) / len(X_test_i)

  print(f'Accuracy on test dataset:{accuracy:.3f}')
  result = model(torch.from_numpy(cyy).float().to(device)).argmax(axis=-1)
  print(result)
  results.append(result)
  print()

  dp=pd.read_excel('附件.xlsx',2)
dp=dp.fillna(0)
dp.head(5)

cyy=dp.iloc[:,2:16].to_numpy()

results